
#include "pows_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>


namespace optiling {
    const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  PowsTilingData tiling;
  int32_t NUM = 24;
  uint32_t sizeofdatatype;
  uint32_t totalLengthAligned;
  auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  auto socVersion = ascendcPlatform.GetSocVersion();
  uint64_t ub_size;
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
  auto aivNum = ascendcPlatform.GetCoreNum();

  uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();

        //获取输入shape信息
  uint32_t inputNum = context->GetInputShape(1)->GetStorageShape().GetShapeSize(); //输入数量
  uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(1)->GetDataType()); //输入类型
  uint32_t inputLength = inputBytes * inputNum; //输入长度

  auto dt = context->GetInputTensor(0)->GetDataType();
  if(dt == ge::DT_FLOAT16 || dt == ge::DT_BF16){
    sizeofdatatype = 2;
    NUM = 5;
    }
  else{ //DT_FLOAT
    sizeofdatatype = 4;
    NUM = 3;
    }

  uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
  uint32_t tiling_size = ((ub_size) / BLOCK_SIZE / 2) / NUM;
  tiling_size = tiling_size <= 8 ? tiling_size : tiling_size / 8 * 8;

  uint32_t block_size = tiling_size * ALIGN_NUM;
  aivNum = (aivNum < totalLength / block_size) ? aivNum : (totalLength / block_size);
  aivNum = aivNum >= 1 ? aivNum : 1;



  uint32_t core_size = (totalLength / aivNum) / (ALIGN_NUM * 8) * (ALIGN_NUM * 8);
  uint32_t core_remain = totalLength - aivNum * core_size;

  //tiling.set_totalLength(totalLength);
  tiling.set_ALIGN_NUM(ALIGN_NUM);
  tiling.set_tiling_size(tiling_size);
  tiling.set_block_size(block_size);
  //tiling.set_aivNum(aivNum);
  tiling.set_core_size(core_size);
  tiling.set_core_remain(core_remain);

  uint32_t total_length = 0, min_length = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
  for (int i = 0; i < 2; ++i) {
      total_length = std::max<uint32_t>(total_length, context->GetInputShape(i)->GetStorageShape().GetShapeSize());
      min_length = std::min<uint32_t>(min_length, context->GetInputShape(i)->GetStorageShape().GetShapeSize());
  }
  uint32_t x1_length = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
  uint32_t x2_length = context->GetInputShape(1)->GetStorageShape().GetShapeSize();


  if (x1_length !=total_length || x2_length != total_length){
      context->SetTilingKey(2);
  }
  else
  {
      context->SetTilingKey(1);
  }

  int32_t numshapes = context->GetInputShape(0)->GetStorageShape().GetDimNum();
  tiling.set_numshapes(numshapes);
  int32_t shape[128];
  for (int k = 0; k < 2; ++k) {
      int32_t *ss = &shape[k * 64];
      const gert::StorageShape* shape = context->GetInputShape(k);
      for (int i = 0; i < shape->GetStorageShape().GetDimNum(); i++) {
          ss[i] = shape->GetStorageShape().GetDim(i);
      }
  }
  tiling.set_shape(shape);

  int32_t shapefull[64];
  for (int k = 0; k < numshapes; ++k) {
      int32_t *ss = &shape[0];
      int32_t *sf = &shapefull[0];
      int32_t tmp = (ss[k] > ss[k + 64]) ? ss[k] : ss[k + 64];   
      sf[k] = tmp;
  }
  tiling.set_shapefull(shapefull);

  tiling.set_total_length(total_length);
  tiling.set_x1_length(x1_length);
  tiling.set_x2_length(x2_length);

  std::cout<< "total_length:" << total_length << std::endl;
  std::cout<< "x1_length:" << x1_length << std::endl;
  std::cout<< "x2_length:" << x2_length << std::endl;
  std::cout<< "numshapes:" << numshapes << std::endl;

  context->SetBlockDim(1);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  size_t *currentWorkspace = context->GetWorkspaceSizes(1);
  currentWorkspace[0] = 0;

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Pows : public OpDef {
public:
    explicit Pows(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Pows);
}
